import os
from dotenv import load_dotenv
# LangChain v0.3+ 最新兼容导入（无冲突+无报红）
from langchain_community.document_loaders import TextLoader  # 文档加载（社区包）
from langchain_text_splitters import RecursiveCharacterTextSplitter  # 文本分割（独立子包）
# 替换 FAISS 为 Chroma（核心修改1）
from langchain_community.vectorstores import Chroma
from langchain_community.embeddings import ModelScopeEmbeddings  # 嵌入模型（社区包）
from langchain_community.chat_models.tongyi import ChatTongyi
from langchain.chains import RetrievalQA
from langchain_core.prompts import PromptTemplate  # Prompt模板（核心包）
from langchain_community.embeddings.dashscope import DashScopeEmbeddings

# 1. 加载环境变量（阿里云 API Key）
load_dotenv()
dashscope_api_key = os.getenv("DASHSCOPE_API_KEY")
if not dashscope_api_key:
    raise ValueError("请在 .env 文件中设置 DASHSCOPE_API_KEY")

# 2. 准备本地知识库（以文本文件为例，可替换为 PDF/Word 等）
knowledge_base_path = "ali_rag_kb.txt"
# 生成示例知识库（阿里云相关知识，可替换为你的业务文档）
if not os.path.exists(knowledge_base_path):
    with open(knowledge_base_path, "w", encoding="utf-8") as f:
        f.write("""
阿里云通义千问（Qwen）是阿里巴巴研发的大语言模型，支持中文、英文等多语言交互，核心特性包括：
1. 中文理解能力强：针对中文语境优化，支持方言、古诗词、专业术语理解；
2. 多任务能力：可实现文本生成、问答、翻译、代码编写、逻辑推理等；
3. API 调用灵活：提供 qwen-turbo（快速版）、qwen-plus（增强版）、qwen-max（旗舰版）等多个版本；
4. 免费额度：新用户注册后可获得一定量的免费调用额度，适合开发测试；
5. 部署方式：支持 API 调用、私有化部署、阿里云百炼平台一键部署。

阿里云 ModelScope 是开源模型社区，提供：
1. 海量开源模型：涵盖 NLP、CV、语音、多模态等领域；
2. 一键调用：支持模型快速加载和推理，无需复杂配置；
3. 中文优化模型：众多模型针对中文场景优化，适合国内开发者；
4. 工具链集成：可与 LangChain、FAISS 等工具无缝配合。

通义千问 API 调用步骤：
1. 登录阿里云百炼平台，创建 API-KEY；
2. 安装 dashscope 依赖包（pip install dashscope）；
3. 配置 API-KEY，调用对应模型版本；
4. 传入 Prompt，获取生成结果。
""")

# 3. 加载文档 + 文本分割（解决长文本处理问题）
loader = TextLoader(knowledge_base_path, encoding="utf-8")
documents = loader.load()  # 加载文档
# 文本分割：将长文档拆分为短文本块（适合嵌入和检索）
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=300,  # 每个文本块最大长度
    chunk_overlap=50,  # 文本块重叠长度（保证上下文连贯）
    length_function=len
)
text_chunks = text_splitter.split_documents(documents)

# 4. 初始化嵌入模型（文本转向量）
# 选用 DashScope 嵌入模型（也可切换回 ModelScope 本地模型）
embeddings = DashScopeEmbeddings(
    model="text-embedding-v1",
    # api_key=dashscope_api_key  # Chroma 需显式传入 API Key（核心修改2）
)
# 可选：切换为 ModelScope 本地嵌入模型（无需 API 调用）
# embeddings = ModelScopeEmbeddings(
#     model_id="iic/text2vec-large-chinese",
#     model_revision="v1.0.0"
# )

# 5. 构建 Chroma 向量数据库（核心修改3：替换 FAISS 为 Chroma）
chroma_db_path = "ali_chroma_db"  # Chroma 持久化目录（自动创建）

# 检查是否已存在 Chroma 数据库，存在则加载，不存在则新建
if os.path.exists(chroma_db_path):
    # 加载已有 Chroma 数据库
    vector_db = Chroma(
        persist_directory=chroma_db_path,
        embedding_function=embeddings
    )
    print(f"✅ 加载已存在的 Chroma 数据库：{chroma_db_path}")
else:
    # 新建 Chroma 数据库并添加文档
    vector_db = Chroma.from_documents(
        documents=text_chunks,
        embedding=embeddings,
        persist_directory=chroma_db_path  # 指定持久化目录（自动保存）
    )
    print(f"✅ 新建 Chroma 数据库并持久化到：{chroma_db_path}")

# 6. 创建检索器（从向量库中召回相关知识）
retriever = vector_db.as_retriever(
    search_kwargs={"k": 2}  # 每次检索返回最相关的 2 个文本块（参数与 FAISS 一致）
)

# 7. 初始化阿里云通义千问 LLM
llm = ChatTongyi(
    model="qwen3-vl-8b-instruct",  # 快速版（免费额度充足，适合测试）
    api_key=dashscope_api_key,
    temperature=0.1,  # 温度越低，回答越确定（避免幻觉）
    top_p=0.8
)

# 8. 自定义 Prompt（引导 LLM 基于检索到的知识回答，避免编造）
prompt_template = """
请严格基于以下参考知识回答用户问题，不要添加任何参考知识之外的信息。
如果参考知识中没有相关答案，直接回复“暂无相关信息”，不要编造内容。

参考知识：
{context}

用户问题：{question}

回答：
"""
prompt = PromptTemplate(
    template=prompt_template,
    input_variables=["context", "question"]
)

# 9. 构建 RAG 链（检索 + 生成）
rag_chain = RetrievalQA.from_chain_type(
    llm=llm,
    chain_type="stuff",  # 直接拼接检索结果给 LLM（适合短文本）
    retriever=retriever,
    chain_type_kwargs={"prompt": prompt},  # 传入自定义 Prompt
    return_source_documents=True  # 返回检索到的源文档（便于验证）
)

# 10. 测试 RAG 问答
def rag_qa(query: str):
    """执行 RAG 问答并输出结果"""
    result = rag_chain.invoke({"query": query})
    answer = result["result"]
    sources = [doc.page_content for doc in result["source_documents"]]

    print(f"📌 问题：{query}")
    print(f"🤖 回答：{answer}")
    print(f"📚 参考来源（共 {len(sources)} 条）：")
    for i, source in enumerate(sources, 1):
        print(f"  {i}. {source.strip()}")
    print("-" * 80)

# 运行测试
if __name__ == "__main__":
    print("=== ModelScope + 阿里云通义千问 + Chroma RAG 演示 ===\n")
    test_queries = [
        "通义千问有哪些版本？",
        "如何调用通义千问 API？",
        "ModelScope 是什么？",
        "通义千问支持私有化部署吗？",
        "阿里云有免费的 LLM 可以用吗？"
    ]
    for query in test_queries:
        rag_qa(query)